#!/usr/bin/env python3
"""
AutoGen 多智能体对话示例 - 使用本地deepseek-coder:1.3b 模型
演示用户代理、助手代理和代码执行代理的协作
"""

import autogen
from autogen import AssistantAgent, UserProxyAgent
from autogen.coding import LocalCommandLineCodeExecutor


def main():
    """主函数：设置和运行多智能体对话"""

    # =============== 关键修改：配置本地 DeepSeek 模型 ===============
    config_list = [
        {
            # 使用 Ollama 提供的本地 API
            "model": "deepseek-coder:1.3b",  # 修改：使用正确的模型名称
            "base_url": "http://localhost:11434/v1",  # 修改：正确的Ollama API地址
            "api_key": "ollama",  # Ollama 不需要真实 API key
        }
    ]

    # LLM配置 - 针对本地模型优化参数
    llm_config = {
        "config_list": config_list,
        "temperature": 0.3,
        "timeout": 180,
        "max_tokens": 2048,
    }

    # 2. 创建代码执行器
    code_executor = LocalCommandLineCodeExecutor(
        timeout=120,
        work_dir="./autogen_workspace",
    )

    # 3. 定义智能体 - 针对小模型优化 system_message

    # 用户代理
    user_proxy = UserProxyAgent(
        name="UserProxy",
        system_message="""你是用户代理。职责：
        1. 代表用户发起任务
        2. 执行代码验证结果
        3. 任务完成时说'TERMINATE'
        
        请保持简洁高效。""",
        human_input_mode="NEVER",
        max_consecutive_auto_reply=8,
        is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE"),
        code_execution_config={
            "executor": code_executor,
        },
        llm_config=llm_config,
    )

    # 数据分析师代理 - 简化角色描述
    data_analyst = AssistantAgent(
        name="DataAnalyst",
        system_message="""你是数据分析师。职责：
        1. 数据分析和可视化
        2. 编写Python代码（pandas、matplotlib）
        3. 提供简洁的数据洞察
        
        代码要求：清晰、有注释、可运行。""",
        llm_config=llm_config,
    )

    # 代码审查员代理 - 简化职责
    code_reviewer = AssistantAgent(
        name="CodeReviewer",
        system_message="""你是代码审查员。职责：
        1. 检查代码正确性
        2. 提出改进建议
        3. 确保代码可运行
        
        重点关注：逻辑正确性、基本优化、错误处理。""",
        llm_config=llm_config,
    )

    # 项目经理代理 - 简化协调功能
    project_manager = AssistantAgent(
        name="ProjectManager",
        system_message="""你是项目经理。职责：
        1. 分解任务
        2. 协调团队工作
        3. 提供最终总结
        
        保持沟通简洁有效。""",
        llm_config=llm_config,
    )

    # 4. 创建群聊
    group_chat = autogen.GroupChat(
        agents=[user_proxy, data_analyst, code_reviewer, project_manager],
        messages=[],
        max_round=15,
        speaker_selection_method="round_robin",
    )

    # 群聊管理器
    manager = autogen.GroupChatManager(
        groupchat=group_chat,
        llm_config=llm_config
    )

    # 5. 定义任务 - 简化任务描述
    task = """
    创建销售数据分析项目：
    
    1. 生成模拟数据：
       - 时间：2023年1-12月
       - 列：日期、产品类别、销售额、销售员、地区
       - 产品类别：电子、服装、食品、书籍
       - 销售员：Alice、Bob、Charlie、Diana
    
    2. 分析：
       - 月度趋势
       - 类别占比
       - 销售员排名
       - 地区分布
    
    3. 创建可视化图表
    4. 提供业务建议
    
    要求：代码可直接运行，包含注释。
"""

    # 6. 启动对话
    print("🚀 启动AutoGen+DeepSeek多智能体协作...")
    print("模型：DeepSeek 1.5B (本地)")
    print("=" * 60)

    user_proxy.initiate_chat(
        manager,
        message=task,
    )

    print("=" * 60)
    print("✅ 任务完成！")


def setup_environment():
    """设置环境变量和工作目录"""
    import os
    import requests

    # 检查 Ollama 服务
    try:
        response = requests.get("http://localhost:11434/api/tags", timeout=5)
        if response.status_code != 200:
            print("⚠️  警告: Ollama 服务未启动")
            print("请先启动 Ollama 服务:")
            print("ollama serve")
            return False

        # 检查模型是否已下载
        models = response.json().get("models", [])
        model_names = [model.get("name", "") for model in models]

        if "deepseek-coder:1.3b" not in model_names:
            print("⚠️  警告: deepseek-coder:1.3b 模型未下载")
            print("请下载模型:")
            print("ollama pull deepseek-coder:1.3b")
            return False

    except Exception as e:
        print(f"⚠️  警告: 无法连接到 Ollama 服务: {e}")
        print("请确保 Ollama 已安装并启动:")
        print("1. 安装: curl -fsSL https://ollama.com/install.sh | sh")
        print("2. 下载模型: ollama pull deepseek-coder:1.3b")
        print("3. 启动服务: ollama serve")
        return False

    # 创建工作目录
    work_dir = "./autogen_workspace"
    if not os.path.exists(work_dir):
        os.makedirs(work_dir)
        print(f"📁 创建工作目录: {work_dir}")

    return True


if __name__ == "__main__":
    print("🤖 AutoGen + DeepSeek 1.5B 多智能体对话示例")
    print("Python 3.12.7 环境 + 本地模型")
    print("=" * 60)

    # 设置环境
    if setup_environment():
        try:
            main()
        except KeyboardInterrupt:
            print("\n👋 用户中断，程序退出")
        except Exception as e:
            print(f"❌ 发生错误: {e}")
            import traceback

            traceback.print_exc()
    else:
        print("❌ 环境设置失败，请检查 Ollama 配置")